#!/usr/bin/env python

# (c) 2007-2011 Helsinki University of Technology and University of Helsinki
# Licensed under the MIT license; see LICENSE.txt for more information.

from worker import Worker
from rdf import *

import copy
import urlparse
from lxml.etree import Element, SubElement, ElementTree

class ContentExtractor (Worker):
  def __init__(self, session):
    Worker.__init__(self, session.getFatalEvent())
    self._session = session
    
    outputdir = self._session.getConfig().getOutputDir()
    source = self._session.getSource().getName()
    self._outputdir = outputdir + "/" + source
  
  def process(self, doc):
    # do sth
    tree = doc.getParseTree()
    
    # find headline
    headline = None
    headlineXPath = self._session.getSource().getHeadlineXPath()
    if headlineXPath:
      elems = tree.xpath(headlineXPath)
      if elems:
        headline = elems[0]
      else:
        doc.log('WARN', 'content', 'Headline extraction with configured XPath expression failed', headlineXPath)
#    else:
#      doc.log('NOTE', 'content', 'No headline XPath specified, not extracting document headline')
        
    # find content
    content = None
    contentXPath = self._session.getSource().getContentXPath()
    if contentXPath:
      elems = tree.xpath(contentXPath)
      if elems:
        content = elems[0]
      else:
        doc.log('WARN', 'content', 'Content extraction with configured XPath expression failed', headlineXPath)
#    else:
#      doc.log('NOTE', 'content', 'No content XPath specified, not extracting document content')

    if self._session.getStore() and headline is not None and content is not None:
      # build the HTML document to store
      html = Element('html')
      head = SubElement(html, 'head')
      title = SubElement(head, 'title')
      # copy the headline to the TITLE element
      title.text = headline.text
      
      # transmogrify the content into a BODY element
      body = copy.deepcopy(content)
      body.tag = 'body'			# switch the tag
      for key in body.attrib.keys():	# clear all attributes
        del body.attrib[key]
      
      # figure out the base URL
      base = tree.find('base')
      if base is not None:
        base = base.attrib.get('href', None)
      if base is None:
        base = doc.getURL()
      
      # make links and image URLs absolute
      for a in body.getiterator('a'):
        href = a.attrib.get('href', None)
        if href is None: continue	# <a name> or whatever
        a.attrib['href'] = urlparse.urljoin(base, href)
      for img in body.getiterator('img'):
        src = img.attrib.get('src', None)
        if src is None: continue	# broken image
        img.attrib['src'] = urlparse.urljoin(base, src)
      
      # move the mogrified content to new doc
      html.append(body)
      
      # create the resulting minimal HTML document
      outtree = ElementTree(html)
      
      outfilename = self._outputdir + "/" + doc.getIdentifierMD5() + ".html"
      outfile = file(outfilename, 'w')
      outtree.write(outfile, 'UTF-8')
      outfile.close()
    
    doc.setCompleted('content')
    self._session.schedule(doc)
    return True
    